#define DTYPE_X half
#define DTYPE_Y half
#define DTYPE_Z half
#include "add_vector_tiling.h"
#include "kernel_operator.h"
constexpr int32_t BUFFER_NUM = 2;

class KernelAddVector {
    public:
        __aicore__ inline KernelAddVector() {}
        __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t smallCoreDataNum,
                                uint32_t bigCoreDataNum, uint32_t finalBigTileNum, 
                                uint32_t finalSmallTileNum, uint32_t tileDataNum, 
                                uint32_t smallTailDataNum, uint32_t bigTailDataNum, 
                                uint32_t tailBlockNum) {
            if (0 == AscendC::GetBlockIdx()) {
                AscendC::printf("bigCoreDataNum=%u, smallCoreDataNum=%u, finalBigTileNum=%u, finalSmallTileNum=%u, tileDataNum=%u, smallTailDataNum=%u, bigTailDataNum=%u, tailBlockNum=%u\n",
                                bigCoreDataNum, smallCoreDataNum, finalBigTileNum, finalSmallTileNum, tileDataNum, smallTailDataNum, bigTailDataNum, tailBlockNum);
            }

            if (0 == AscendC::GetBlockIdx() || tailBlockNum == AscendC::GetBlockIdx()) {
                AscendC::printf("Init start: block number=%u, blockIdx=%u\n", AscendC::GetBlockNum(), AscendC::GetBlockIdx());
            }

            uint32_t coreNum = AscendC::GetBlockIdx();
            uint32_t globalBufferIndex = bigCoreDataNum * AscendC::GetBlockIdx();
            this->tileDataNum = tileDataNum;
            if (coreNum < tailBlockNum) { 
                this->coreDataNum = bigCoreDataNum;
                this->tileNum = finalBigTileNum;
                this->tailDataNum = bigTailDataNum;
            }
            else { 
                this->coreDataNum = smallCoreDataNum;
                this->tileNum = finalSmallTileNum;
                this->tailDataNum = smallTailDataNum;
                globalBufferIndex -= (bigCoreDataNum - smallCoreDataNum) * (AscendC::GetBlockIdx() - tailBlockNum);
            }
            xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x + globalBufferIndex, this->coreDataNum);
            yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y + globalBufferIndex, this->coreDataNum);
            zGm.SetGlobalBuffer((__gm__ DTYPE_Z*)z + globalBufferIndex, this->coreDataNum);
            pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X));
            pipe.InitBuffer(inQueueY, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_Y));
            pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_Z));
            if (0 == AscendC::GetBlockIdx() || tailBlockNum == AscendC::GetBlockIdx()) {
                AscendC::printf("Init end: (SetGlobalBuffer) globalBufferIndex=%u, coreDataNum=%u, (InitBuffer for Queue) tileDataNum=%u\n", globalBufferIndex, this->coreDataNum, this->tileDataNum);
            }
        }

        __aicore__ inline void Process() {
            int32_t loopCount = this->tileNum;
            this->processDataNum = this->tileDataNum;
            for (int32_t i = 0; i < loopCount; i++) {
                if (i == this->tileNum - 1) {
                    this->processDataNum = this->tailDataNum;
                }
                if (0 == AscendC::GetBlockIdx()) {
                    AscendC::printf("范式计算第 %u 回: blockIdx=%u, processDataNum=%u\n", i+1, AscendC::GetBlockIdx(), this->processDataNum);
                }
                CopyIn(i);
                Compute(i);
                CopyOut(i);
            }
        }

    private:
        __aicore__ inline void CopyIn(int32_t progress) {
            AscendC::LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
            AscendC::LocalTensor<DTYPE_Y> yLocal = inQueueY.AllocTensor<DTYPE_Y>();
            AscendC::DataCopy(xLocal, xGm[progress * this->tileDataNum], this->processDataNum);
            AscendC::DataCopy(yLocal, yGm[progress * this->tileDataNum], this->processDataNum);
            inQueueX.EnQue(xLocal);
            inQueueY.EnQue(yLocal);
        }

        __aicore__ inline void Compute(int32_t progress) {
            AscendC::LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
            AscendC::LocalTensor<DTYPE_Y> yLocal = inQueueY.DeQue<DTYPE_Y>();
            AscendC::LocalTensor<DTYPE_Z> zLocal = outQueueZ.AllocTensor<DTYPE_Z>();
            AscendC::Add(zLocal, xLocal, yLocal, this->processDataNum);
            outQueueZ.EnQue<DTYPE_Z>(zLocal);
            inQueueX.FreeTensor(xLocal);
            inQueueY.FreeTensor(yLocal);

        }

        __aicore__ inline void CopyOut(int32_t progress) {
            AscendC::LocalTensor<DTYPE_Z> zLocal = outQueueZ.DeQue<DTYPE_Z>();
            AscendC::DataCopy(zGm[progress * this->tileDataNum], zLocal, this->processDataNum);
            outQueueZ.FreeTensor(zLocal);
        }

    private:
        AscendC::TPipe pipe;
        AscendC::TQue<AscendC::QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueY;
        AscendC::TQue<AscendC::QuePosition::VECOUT, BUFFER_NUM> outQueueZ;
        AscendC::GlobalTensor<DTYPE_X> xGm;
        AscendC::GlobalTensor<DTYPE_Y> yGm;
        AscendC::GlobalTensor<DTYPE_Z> zGm;
        uint32_t coreDataNum;
        uint32_t tileNum;
        uint32_t tileDataNum;
        uint32_t tailDataNum;
        uint32_t processDataNum;
};

extern "C" __global__ __aicore__ void add_vector(GM_ADDR x, GM_ADDR y, GM_ADDR z, AddVectorTilingData tiling) {
    KernelAddVector op;
    op.Init(x, y, z, tiling.smallCoreDataNum, tiling.bigCoreDataNum, tiling.finalBigTileNum, tiling.finalSmallTileNum, tiling.tileDataNum, tiling.smallTailDataNum, tiling.bigTailDataNum, tiling.tailBlockNum);
    op.Process();
}
